import os
import time
from channels.channels import Channel

class Lines2Split(Channel):
    """Convert lines to train, valid, test"""
    required = ['lines']
    send = ['lines']
    
    @staticmethod
    def execute(args={}, meta={}, lines=[], **kwargs):
        split = args.split
        # check split
        if len(split) != 3 or sum(split) != 10:
            raise Exception('The param split have to have three number and the sum is equal to 10;')

        lines_len = len(lines)
        train_nums = int(lines_len * split[0] / 10)
        valid_nums = int(lines_len * split[1] / 10)
        test_nums = lines_len - train_nums - valid_nums

        cur_time = time.strftime('_%Y年_%m月_%d日%H时%M分%S秒',time.localtime(time.time()))
        save_dir_name = os.path.basename(args.save_path).split('.')[0] + '_' + cur_time + '__' + '_'.join([str(train_nums), str(valid_nums), str(test_nums)])      
        save_dir = os.path.join(args.split_dir, save_dir_name)
        if not os.path.isdir(save_dir):
            os.makedirs(save_dir)


        train_file = os.path.join(save_dir, 'train.txt')
        valid_file = os.path.join(save_dir, 'valid.txt')
        test_file = os.path.join(save_dir, 'test.txt')


        with open(train_file, 'w', encoding='utf-8') as fw:
            print('train_file saving:')
            for index, line in enumerate(lines[: train_nums]):
                # bar show
                args.bar_show(index, train_nums)

                fw.write(line + '\n')

        with open(valid_file, 'w', encoding='utf-8') as fw:
            print('valid_file saving:')
            for index, line in enumerate(lines[train_nums: (train_nums + valid_nums)]):
                # bar show
                args.bar_show(index, valid_nums)

                fw.write(line + '\n')
        
        with open(test_file, 'w', encoding='utf-8') as fw:
            print('test_file saving:')
            for index, line in enumerate(lines[(train_nums + valid_nums):]):
                # bar show
                args.bar_show(index, test_nums)

                fw.write(line + '\n')
        
        return {
            'args': args,
            'meta': meta, 
            'lines': lines
        }